#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

template<typename T>
class KernelClipByValue {
 public:
  __aicore__ inline KernelClipByValue() {}
  __aicore__ inline void Init(GM_ADDR x, GM_ADDR clip_value_min, GM_ADDR clip_value_max, GM_ADDR y,
                              uint32_t blockLength,
                              uint32_t tileNum, uint32_t tileLength,
                              uint32_t lasttileLength, uint32_t formerNum,
                              uint32_t formerLength, uint32_t formertileNum,
                              uint32_t formertileLength,
                              uint32_t formerlasttileLength, uint32_t tailNum,
                              uint32_t tailLength, uint32_t tailtileNum,
                              uint32_t tailtileLength,
                              uint32_t taillasttileLength, uint32_t tilingKey, uint32_t datatype) {
    ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
    //this->cof = static_cast<T>(cof);
    this->datatype = datatype;
    clip_value_minGm.SetGlobalBuffer((__gm__ T*)clip_value_min , 1);
    this->clip_value_min = clip_value_minGm.GetValue(0);
    clip_value_maxGm.SetGlobalBuffer((__gm__ T*)clip_value_max , 1);
    this->clip_value_max = clip_value_maxGm.GetValue(0);

    if (tilingKey == 1) {
      this->blockLength = blockLength;
      this->tileNum =
          tileNum ASSERT(tileNum != 0 && "tile num can not be zero!");
      this->tileLength = tileLength / BUFFER_NUM;
      this->lasttileLength = lasttileLength;

      xGm.SetGlobalBuffer((__gm__ T*)x + this->blockLength * GetBlockIdx(),
                          this->blockLength);
      yGm.SetGlobalBuffer(
          (__gm__ T*)y + this->blockLength * GetBlockIdx(),
          this->blockLength);
    }

    if (tilingKey == 2) {
      this->formerNum = formerNum;
      this->formerLength = formerLength;
      this->formertileNum = formertileNum;
      this->formertileLength = formertileLength;
      this->formerlasttileLength = formerlasttileLength;

      this->tailNum = tailNum;
      this->tailLength = tailLength;
      this->tailtileNum = tailtileNum;
      this->tailtileLength = tailtileLength;
      this->taillasttileLength = taillasttileLength;

      if (GetBlockIdx() < this->formerNum) {  //分到大块核的处理
        this->tileLength = this->formertileLength / BUFFER_NUM;
        this->lasttileLength = this->formerlasttileLength;
        this->tileNum = this->formertileNum * BUFFER_NUM;
        xGm.SetGlobalBuffer(
            (__gm__ T*)x + this->formerLength * GetBlockIdx(),
            this->formerLength);
        yGm.SetGlobalBuffer(
            (__gm__ T*)y + this->formerLength * GetBlockIdx(),
            this->formerLength);
      } else {  //分到小块核的处理，需要处理的数据量比大核少alignNum个
        this->tileLength = this->tailtileLength / BUFFER_NUM;
        this->lasttileLength = this->taillasttileLength;
        this->tileNum = this->tailtileNum * BUFFER_NUM;
        xGm.SetGlobalBuffer(
            (__gm__ T*)x + this->formerLength * this->formerNum +
                this->tailLength * (GetBlockIdx() - this->formerNum),
            this->tailLength);
        yGm.SetGlobalBuffer(
            (__gm__ T*)y + this->formerLength * this->formerNum +
                this->tailLength * (GetBlockIdx() - this->formerNum),
            this->tailLength);
      }
    }

    pipe.InitBuffer(inQueueIN, BUFFER_NUM, this->tileLength * sizeof(T));
    pipe.InitBuffer(outQueueOUT, BUFFER_NUM, this->tileLength * sizeof(T));
  }
  __aicore__ inline void Process() {
    int32_t loopCount = this->tileNum * BUFFER_NUM;
    for (int32_t i = 0; i < loopCount; i++) {
      CopyIn(i);
      Compute(i);
      CopyOut(i);
    }
  }

 private:
  __aicore__ inline void CopyIn(int32_t progress) {
    LocalTensor<T> inLocal = inQueueIN.AllocTensor<T>();

    if (BUFFER_NUM == 1) {
      if (progress == this->tileNum - 1) {
        if (progress == 0) {
          //如果只有一包，则搬运的起始地址为0，tileLength为实际分块的数据量
          DataCopy(inLocal[0], xGm[0], this->tileLength);
        } else {
          //将最后一个分块的起始地址向前移动tileLength-lasttileLength
          DataCopy(
              inLocal[0],
              xGm[(progress - 1) * this->tileLength + this->lasttileLength],
              this->tileLength);
        }
      } else {
        DataCopy(inLocal[0], xGm[progress * this->tileLength],
                 this->tileLength);
      }
    }
    if (BUFFER_NUM == 2) {
      //开启double
      //buffer时，由于将输入数据分成了相等的2部分，分块大小为不开启double
      //buffer的一半， 所以需要对最后两个分块数据的起始地址做处理
      if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1))) {
        //分块大小变为tileLength的一半
        //倒数第2个分块数据的起始地址向前移动（tileLength-lasttileLength)，最后一个分块的起始地址以此为基础进行移动
        DataCopy(
            inLocal[0],
            xGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            (this->tileLength));
      }

      else {
        DataCopy(inLocal[0], xGm[progress * (this->tileLength)],
                 (this->tileLength));
      }
    }

    inQueueIN.EnQue(inLocal);
  }
  __aicore__ inline void Compute(int32_t progress) {
    LocalTensor<T> inLocal = inQueueIN.DeQue<T>();

    LocalTensor<T> yLocal = outQueueOUT.AllocTensor<T>();
    DataCopy(yLocal, inLocal, this->tileLength);
    T value;
    pipe_barrier(PIPE_ALL);

    if(this->datatype == 0 || this->datatype == 1){
      for(int i=0; i<this->tileLength; i++){
        value = inLocal.GetValue(i);
        if((float)value < (float)this->clip_value_min){
            yLocal.SetValue(i,this->clip_value_min);
        }else if((float)value > (float)this->clip_value_max){
            yLocal.SetValue(i,this->clip_value_max);
        }
      }
    }else {//int32需要类型转换，暂时先不考虑
      for(int i=0; i<this->tileLength; i++){
        value = inLocal.GetValue(i);
        if((int)value < (int)this->clip_value_min){
            yLocal.SetValue(i,this->clip_value_min);
        }else if((int)value > (int)this->clip_value_max){
            yLocal.SetValue(i,this->clip_value_max);
        }
      }
    }
    outQueueOUT.EnQue<T>(yLocal);
    inQueueIN.FreeTensor(inLocal);

    // LocalTensor<T> inLocal = inQueueIN.DeQue<T>();

    // LocalTensor<T> yLocal = outQueueOUT.AllocTensor<T>();
    // pipe_barrier(PIPE_ALL);
    // if(this->datatype == 0 || this->datatype == 1){
    //   for(int i=0; i<this->tileLength; i++){
    //     //pipe_barrier(PIPE_ALL);
    //     if((float)inLocal.GetValue(i) < (float)this->clip_value_min){
    //         yLocal.SetValue(i,this->clip_value_min);
    //     }else if((float)inLocal.GetValue(i) > (float)this->clip_value_max){
    //         yLocal.SetValue(i,this->clip_value_max);
    //     }
    //   }
    // }else{
    //   for(int i=0; i<this->tileLength; i++){
    //     //pipe_barrier(PIPE_ALL);
    //     if((int)inLocal.GetValue(i) < (int)this->clip_value_min){
    //         yLocal.SetValue(i,this->clip_value_min);
    //     }else if((int)inLocal.GetValue(i) > (int)this->clip_value_max){
    //         yLocal.SetValue(i,this->clip_value_max);
    //     }
    //   }
    // }
    // // pipe_barrier(PIPE_ALL);
    // outQueueOUT.EnQue<T>(yLocal);

    // inQueueIN.FreeTensor(inLocal);

  }
  __aicore__ inline void CopyOut(int32_t progress) {
    LocalTensor<T> yLocal = outQueueOUT.DeQue<T>();

    if (BUFFER_NUM == 1) {
      if (progress == this->tileNum - 1) {
        if (progress == 0) {
          //如果只有一包，则搬运的起始地址为0，tileLength为实际分块的数据量
          DataCopy(yGm[0], yLocal, this->tileLength);
        } else {
          //将最后一个分块的起始地址向前移动tileLength-lasttileLength
          DataCopy(
              yGm[(progress - 1) * this->tileLength + this->lasttileLength],
              yLocal, this->tileLength);
        }
      } else {
        DataCopy(yGm[progress * this->tileLength], yLocal,
                 this->tileLength);
      }
    }
    if (BUFFER_NUM == 2) {
      //开启double
      //buffer时，由于将输入数据分成了相等的2部分，分块大小为不开启double
      //buffer的一半， 所以需要对最后两个分块数据的起始地址做处理
      if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1))) {
        //分块大小变为tileLength的一半
        //倒数第2个分块数据的起始地址向前移动（tileLength-lasttileLength)，最后一个分块的起始地址以此为基础进行移动
        DataCopy(
            yGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            yLocal, (this->tileLength));
      }

      else {
        DataCopy(yGm[progress * (this->tileLength)], yLocal,
                 (this->tileLength));
      }
    }

    outQueueOUT.FreeTensor(yLocal);
  }

 private:
  TPipe pipe;
  // TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueY, inQueueZ;
  TQue<QuePosition::VECIN, BUFFER_NUM> inQueueIN;
  TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueOUT;
  GlobalTensor<T> xGm;
  GlobalTensor<T> clip_value_minGm, clip_value_maxGm;
  GlobalTensor<T> yGm;
  T clip_value_min;
  T clip_value_max;
  uint32_t datatype;
  uint32_t blockLength;
  uint32_t tileNum;
  uint32_t tileLength;
  uint32_t lasttileLength;
  uint32_t formerNum;
  uint32_t formerLength;
  uint32_t formertileNum;
  uint32_t formertileLength;
  uint32_t formerlasttileLength;
  uint32_t tailNum;
  uint32_t tailLength;
  uint32_t tailtileNum;
  uint32_t tailtileLength;
  uint32_t taillasttileLength;
};

extern "C" __global__ __aicore__ void clip_by_value(GM_ADDR x, GM_ADDR clip_value_min, GM_ADDR clip_value_max, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    if(tiling_data.datatype == 1){
        //half
        KernelClipByValue<half> op;
        uint32_t tilingKey = 1;
        if (TILING_KEY_IS(1)) {
            tilingKey = 1;
        } else if (TILING_KEY_IS(2)) {
            tilingKey = 2;
        } else {
            tilingKey = 1;
        }

        op.Init(x, clip_value_min, clip_value_max, y, tiling_data.blockLength,
                tiling_data.tileNum, tiling_data.tileLength,
                tiling_data.lasttileLength, tiling_data.formerNum,
                tiling_data.formerLength, tiling_data.formertileNum,
                tiling_data.formertileLength, tiling_data.formerlasttileLength,
                tiling_data.tailNum, tiling_data.tailLength, tiling_data.tailtileNum,
                tiling_data.tailtileLength, tiling_data.taillasttileLength,
                tilingKey, tiling_data.datatype);
        op.Process();
    }else if(tiling_data.datatype == 0){
      //float32
        KernelClipByValue<float> op;
        uint32_t tilingKey = 1;
        if (TILING_KEY_IS(1)) {
            tilingKey = 1;
        } else if (TILING_KEY_IS(2)) {
            tilingKey = 2;
        } else {
            tilingKey = 1;
        }

        op.Init(x, clip_value_min, clip_value_max, y, tiling_data.blockLength,
                tiling_data.tileNum, tiling_data.tileLength,
                tiling_data.lasttileLength, tiling_data.formerNum,
                tiling_data.formerLength, tiling_data.formertileNum,
                tiling_data.formertileLength, tiling_data.formerlasttileLength,
                tiling_data.tailNum, tiling_data.tailLength, tiling_data.tailtileNum,
                tiling_data.tailtileLength, tiling_data.taillasttileLength,
                tilingKey, tiling_data.datatype);
        op.Process();
    }else if(tiling_data.datatype == 3){
      //int32
        KernelClipByValue<int32_t> op;
        uint32_t tilingKey = 1;
        if (TILING_KEY_IS(1)) {
            tilingKey = 1;
        } else if (TILING_KEY_IS(2)) {
            tilingKey = 2;
        } else {
            tilingKey = 1;
        }

        op.Init(x, clip_value_min, clip_value_max, y, tiling_data.blockLength,
                tiling_data.tileNum, tiling_data.tileLength,
                tiling_data.lasttileLength, tiling_data.formerNum,
                tiling_data.formerLength, tiling_data.formertileNum,
                tiling_data.formertileLength, tiling_data.formerlasttileLength,
                tiling_data.tailNum, tiling_data.tailLength, tiling_data.tailtileNum,
                tiling_data.tailtileLength, tiling_data.taillasttileLength,
                tilingKey, tiling_data.datatype);
        op.Process();
    }
}